package com.bw.test1;

import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.regression.LinearRegressionModel;
import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.evaluation.EvalRegressionBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.RegressionMetrics;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class EvalRegressionBatchOpTest {
    @Test
    public void testEvalRegressionBatchOp() throws Exception {


        BatchOperator.setParallelism(1);
        // 1. 数据源
        List<Row> df = Arrays.asList(
                Row.of(2, 1, 1),
                Row.of(3, 2, 1),
                Row.of(4, 3, 2),
                Row.of(2, 4, 1),
                Row.of(2, 2, 1),
                Row.of(4, 3, 2),
                Row.of(1, 2, 1)
        );
        // 选择特征和label
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");
        String[] colnames = new String[]{"f0", "f1"};


        //  训练模型
        LinearRegression lr = new LinearRegression()
                .setFeatureCols(colnames)
                .setLabelCol("label")
                .setPredictionCol("pred")
                .enableLazyPrintModelInfo();
        LinearRegressionModel model = lr.fit(batchData);


        // 输出结果
        BatchOperator<?> result = model.transform(batchData);

        result.print();

        // 评估指标
        RegressionMetrics metrics = new EvalRegressionBatchOp().setPredictionCol("pred").setLabelCol("label").linkFrom(
                result).collectMetrics();
        System.out.println("Total Samples Number:" + metrics.getCount());
        System.out.println("SSE:" + metrics.getSse());
        System.out.println("SAE:" + metrics.getSae());
        System.out.println("RMSE:" + metrics.getRmse());
        System.out.println("R2:" + metrics.getR2());
        System.out.println("MSE:" + metrics.getMse());
    }
}